// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/exec/exchange_queue.h>
#include <pollux/exec/exchange_source.h>

namespace kumo::pollux::exec {
    // Handle for a set of producers. This may be shared by multiple Exchanges, one
    // per consumer thread.
    class ExchangeClient : public std::enable_shared_from_this<ExchangeClient> {
    public:
        static constexpr int32_t kDefaultMaxQueuedBytes = 32 << 20; // 32 MB.
        static constexpr std::chrono::milliseconds kRequestDataMaxWait{100};
        static inline const std::string kBackgroundCpuTimeMs = "backgroundCpuTimeMs";

        ExchangeClient(
            std::string taskId,
            int destination,
            int64_t maxQueuedBytes,
            int32_t numberOfConsumers,
            uint64_t minOutputBatchBytes,
            memory::MemoryPool *pool,
            melon::Executor *executor,
            int32_t requestDataSizesMaxWaitSec = 10)
            : taskId_{std::move(taskId)},
              destination_(destination),
              maxQueuedBytes_{maxQueuedBytes},
              kRequestDataSizesMaxWaitSec_{requestDataSizesMaxWaitSec},
              pool_(pool),
              executor_(executor),
              queue_(std::make_shared<ExchangeQueue>(
                  numberOfConsumers,
                  minOutputBatchBytes)),
              // See comment in 'pickSourcesToRequestLocked' for why this is needed
              minOutputBatchBytes_(
                  std::max(static_cast<uint64_t>(1), minOutputBatchBytes)) {
            POLLUX_CHECK_NOT_NULL(pool_);
            POLLUX_CHECK_NOT_NULL(executor_);
            // NOTE: the executor is used to run async response callback from the
            // exchange source. The provided executor must not be
            // melon::InlineLikeExecutor, otherwise it might cause potential deadlock as
            // the response callback in exchange client might call back into the
            // exchange source under uncertain execution context. For instance, the
            // exchange client might inline close the exchange source from a background
            // thread of the exchange source, and the close needs to wait for this
            // background thread to complete first.
            POLLUX_CHECK_NULL(dynamic_cast<const melon::InlineLikeExecutor*>(executor_));
            POLLUX_CHECK_GE(
                destination, 0, "Exchange client destination must not be negative");
        }

        ~ExchangeClient();

        memory::MemoryPool *pool() const {
            return pool_;
        }

        // Creates an exchange source and starts fetching data from the specified
        // upstream task. If 'close' has been called already, creates an exchange
        // source and immediately closes it to notify the upstream task that data is
        // no longer needed. Repeated calls with the same 'taskId' are ignored.
        void addRemoteTaskId(const std::string &remoteTaskId);

        void noMoreRemoteTasks();

        // Closes exchange sources.
        void close();

        // Returns runtime statistics aggregated across all of the exchange sources.
        // ExchangeClient is expected to report background CPU time by including a
        // runtime metric named ExchangeClient::kBackgroundCpuTimeMs.
        melon::F14FastMap<std::string, RuntimeMetric> stats() const;

        const std::shared_ptr<ExchangeQueue> &queue() const {
            return queue_;
        }

        /// Returns up to 'maxBytes' pages of data, but no less than one.
        ///
        /// If no data is available returns empty list and sets 'atEnd' to true if no
        /// more data is expected. If data is still expected, sets 'atEnd' to false
        /// and sets 'future' to a Future that will complete when data arrives.
        ///
        /// The data may be compressed, in which case 'maxBytes' applies to compressed
        /// size.
        std::vector<std::unique_ptr<SerializedPage> >
        next(int consumerId, uint32_t maxBytes, bool *atEnd, ContinueFuture *future);

        std::string toString() const;

        melon::Dynamic toJson() const;

        std::chrono::seconds requestDataSizesMaxWaitSec() const {
            return kRequestDataSizesMaxWaitSec_;
        }

    private:
        struct RequestSpec {
            std::shared_ptr<ExchangeSource> source;

            // How much bytes to request from this source.  0 bytes means request data
            // sizes only.
            int64_t maxBytes;
        };

        struct ProducingSource {
            std::shared_ptr<ExchangeSource> source;
            std::vector<int64_t> remainingBytes;
        };

        std::vector<RequestSpec> pickSourcesToRequestLocked();

        void request(std::vector<RequestSpec> &&requestSpecs);

        // Handy for ad-hoc logging.
        const std::string taskId_;
        const int destination_;
        const int64_t maxQueuedBytes_;
        const std::chrono::seconds kRequestDataSizesMaxWaitSec_;

        memory::MemoryPool *const pool_;
        melon::Executor *const executor_;
        const std::shared_ptr<ExchangeQueue> queue_;

        std::unordered_set<std::string> remoteTaskIds_;
        std::vector<std::shared_ptr<ExchangeSource> > sources_;
        bool closed_{false};

        // The minimum byte size the consumer is expected to consume from
        // the exchange queue.
        const uint64_t minOutputBatchBytes_;

        // Total number of bytes in flight.
        int64_t totalPendingBytes_{0};

        // A queue of sources that have returned non-empty response from the latest
        // request.
        std::queue<ProducingSource> producingSources_;
        // A queue of sources that returned empty response from the latest request.
        std::queue<std::shared_ptr<ExchangeSource> > emptySources_;
    };
} // namespace kumo::pollux::exec
